from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser

from com.wp.langchain.exercise.LLMInit import LLMInit
llm = LLMInit().get_llm();
# 直接提供问题，并调用llm
# 创建提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "您是世界级的技术文档编写者。"),
    ("user", "{input}")
])

# 使用输出解析器
#output_parser = StrOutputParser()
output_parser = JsonOutputParser()

# 创建链
chain = prompt | llm | output_parser
# 调用链并获取响应
# 调用它并提出同样的问题。答案是一个字符串，而不是ChatMessage
# chain.invoke({"input": "LangChain是什么?"})
response = chain.invoke({"input": "LangChain是什么? 问题用question 回答用answer 用JSON格式回复"})
print(response)